#include "../src/estep.hh"
#include "../src/nominal.hh"

#include <cassert>
#include <iostream>

namespace
{
using namespace playerpiano;

struct EstepTestDatum 
{ 
  struct
    {
      float    priorz[5];
      Rating   ratings[8];
      float    alpha[4 * 5 * 5];
      float    logbeta;
    }                   inputs;

  float pz[5];
};

static void
test_estep (void)
{
  EstepTestDatum data[] = 
    {{{{0.365583, 0.524757, 0.862708, 0.319783,  0.26922}, {{2, 4}, {3, 2}, {1, 2}, {3, 2}, {2, 2}, {1, 1}, {3,  2}, {1, 1}}, {-1., 0.486307, 0.342955, 0.128267, 0.871339,  0.383534, -1., 0.235787, 0.831097, 0.475512, 0.284733,  0.514043, -1., 0.451854, 0.623866, 0.306703, 0.87594,  0.34407, -1., 0.86006, 0.612388, 0.374104, 0.209263,  0.494477, -1., -1., 0.087631, 0.511396, 0.88948, 0.225257,  0.601324, -1., 0.168441, 0.761213, 0.353919, 0.21779,  0.932654, -1., 0.930116, 0.878406, 0.933058, 0.418611,  0.478262, -1., 0.25454, 0.626355, 0.542671, 0.134192,  0.39448, -1., -1., 0.0139672, 0.168568, 0.924929, 0.900003,  0.926336, -1., 0.657172, 0.0354497, 0.674746, 0.325012,  0.488731, -1., 0.274237, 0.320827, 0.107222, 0.556077,  0.34412, -1., 0.442421, 0.174164, 0.137466, 0.865858,  0.187881, -1., -1., 0.547809, 0.594794, 0.731666, 0.7934,  0.533842, -1., 0.426227, 0.806737, 0.893397, 0.607506,  0.769054, -1., 0.771287, 0.218652, 0.282494, 0.280323,  0.49705, -1., 0.897825, 0.175272, 0.724246, 0.15293,  0.455404, -1.}, 0.00110741}, {0.000573876, 0.0154102, 0.981777,  0.000394054,  0.00184521}}, {{{0.58678, 0.287072, 0.267524, 0.453298,  0.991986}, {{3, 1}, {2, 0}, {0, 4}, {1, 1}, {2, 2}, {0, 3}, {2,  3}, {0, 3}}, {-1., 0.555406, 0.474123, 0.919456, 0.565759,  0.748669, -1., 0.580726, 0.31195, 0.796705, 0.977382,  0.362074, -1., 0.029456, 0.516382, 0.480332, 0.464249,  0.854184, -1., 0.792136, 0.327402, 0.00884494, 0.853077,  0.205356, -1., -1., 0.0403305, 0.741321, 0.399779, 0.21337,  0.484925, -1., 0.267198, 0.480323, 0.647611, 0.736256,  0.686472, -1., 0.168373, 0.850906, 0.758874, 0.324398,  0.138917, -1., 0.334524, 0.278542, 0.860149, 0.284733,  0.542388, -1., -1., 0.951139, 0.851304, 0.431656, 0.337033,  0.910809, -1., 0.109982, 0.0318773, 0.123663, 0.425884,  0.842785, -1., 0.551554, 0.476052, 0.689628, 0.156313,  0.383181, -1., 0.625146, 0.930755, 0.831915, 0.244264,  0.290622, -1., -1., 0.652213, 0.971766, 0.959531, 0.748234,  0.701074, -1., 0.120463, 0.527875, 0.411201, 0.790265,  0.0104801, -1., 0.495998, 0.287538, 0.364381, 0.167696,  0.944443, -1., 0.811486, 0.674752, 0.0113829, 0.561262,  0.18634, -1.}, 0.743997}, {0.000310669, 0.0417779, 0.00927661,  0.939032,  0.00960276}}, {{{0.179468, 0.316998, 0.895718, 0.091784,  0.207702}, {{2, 1}, {1, 0}, {3, 3}, {3, 2}, {3, 3}, {0, 4}, {1,  2}, {2, 1}}, {-1., 0.357467, 0.147485, 0.39071, 0.0872391,  0.829592, -1., 0.736284, 0.600445, 0.076759, 0.333594,  0.448746, -1., 0.236064, 0.909063, 0.389151, 0.637259,  0.561312, -1., 0.897681, 0.827889, 0.450919, 0.817314,  0.718213, -1., -1., 0.510891, 0.555201, 0.72553, 0.510511,  0.153424, -1., 0.407716, 0.33482, 0.423272, 0.323832,  0.671432, -1., 0.734375, 0.346513, 0.990237, 0.222687,  0.498311, -1., 0.437449, 0.601086, 0.585427, 0.937,  0.539769, -1., -1., 0.773197, 0.134508, 0.119685, 0.821556,  0.262306, -1., 0.579308, 0.394155, 0.311045, 0.108882,  0.171592, -1., 0.0593345, 0.887774, 0.785051, 0.50016,  0.324959, -1., 0.541261, 0.794813, 0.277474, 0.826648,  0.103811, -1., -1., 0.193727, 0.692046, 0.889648, 0.564043,  0.42053, -1., 0.557538, 0.769963, 0.742487, 0.158224,  0.97823, -1., 0.375808, 0.431441, 0.0493414, 0.806638,  0.316473, -1., 0.543668, 0.264291, 0.306478, 0.991514,  0.00240712, -1.}, 0.469477}, {0.000986406, 0.184973, 0.788615,  0.0115913,  0.0138347}}, {{{0.0290043, 0.164867, 0.898596, 0.27575,  0.336958}, {{3, 4}, {0, 4}, {1, 2}, {0, 3}, {0, 4}, {2, 0}, {2,  0}, {0, 0}}, {-1., 0.275218, 0.334553, 0.85522, 0.77942,  0.505256, -1., 0.592066, 0.696996, 0.80119, 0.129448,  0.160625, -1., 0.647655, 0.994553, 0.812974, 0.616957,  0.383364, -1., 0.688075, 0.82146, 0.61455, 0.913887,  0.65907, -1., -1., 0.656593, 0.715954, 0.638137, 0.322112,  0.381375, -1., 0.381401, 0.782917, 0.542692, 0.876119,  0.789335, -1., 0.0859205, 0.741502, 0.746672, 0.62871,  0.438266, -1., 0.746949, 0.933697, 0.0117531, 0.0549014,  0.0588744, -1., -1., 0.112237, 0.397203, 0.141015, 0.399804,  0.455644, -1., 0.681249, 0.502878, 0.0776916, 0.0742691,  0.299848, -1., 0.719961, 0.534999, 0.19815, 0.510513,  0.634041, -1., 0.793498, 0.451478, 0.881803, 0.195775,  0.0465487, -1., -1., 0.517781, 0.870049, 0.140874, 0.987674,  0.405544, -1., 0.472846, 0.999859, 0.58787, 0.9499, 0.791597, -1.,  0.496981, 0.510179, 0.87563, 0.49175, 0.77702, -1., 0.975179,  0.677481, 0.981237, 0.14298, 0.181681, -1.},  0.226002}, {0.00790623, 0.000735052, 0.0453077, 0.0121512,  0.9339}}, {{{0.0994343, 0.947205, 0.135133, 0.708221,  0.229385}, {{1, 3}, {0, 4}, {0, 3}, {1, 4}, {1, 1}, {1, 4}, {0,  2}, {1, 1}}, {-1., 0.806331, 0.147458, 0.302678, 0.756539,  0.806472, -1., 0.559588, 0.352778, 0.964941, 0.309491,  0.0494093, -1., 0.477148, 0.473192, 0.532471, 0.07423,  0.799667, -1., 0.491955, 0.389491, 0.892549, 0.573665,  0.392521, -1., -1., 0.442287, 0.757416, 0.865443, 0.163136,  0.635956, -1., 0.609957, 0.562766, 0.406597, 0.829484,  0.0503692, -1., 0.209987, 0.441656, 0.519993, 0.00095985,  0.732839, -1., 0.968464, 0.987522, 0.92673, 0.933172,  0.476509, -1., -1., 0.598031, 0.0341812, 0.359507, 0.0839889,  0.155744, -1., 0.276765, 0.494064, 0.920853, 0.519788,  0.666808, -1., 0.931298, 0.514256, 0.690305, 0.616439,  0.721311, -1., 0.0725998, 0.170312, 0.615479, 0.988472,  0.104136, -1., -1., 0.18279, 0.688749, 0.0552995, 0.627626,  0.584759, -1., 0.654568, 0.695792, 0.543637, 0.429015,  0.377803, -1., 0.201728, 0.622784, 0.909227, 0.710995,  0.27043, -1., 0.108528, 0.218922, 0.0945556, 0.549119,  0.0359285, -1.}, 0.0486105}, {0.00350204, 0.390004, 0.0231149,  0.193114,  0.390265}}, {{{0.479077, 0.560648, 0.931793, 0.865821,  0.790327}, {{3, 4}, {2, 4}, {2, 2}, {0, 2}, {1, 3}, {2, 4}, {0,  4}, {1, 0}}, {-1., 0.505348, 0.304167, 0.281062, 0.135759,  0.809556, -1., 0.76053, 0.852046, 0.757957, 0.607827,  0.137746, -1., 0.94282, 0.0469621, 0.337397, 0.0292174,  0.723897, -1., 0.952407, 0.788278, 0.993289, 0.675287,  0.47333, -1., -1., 0.22763, 0.0614959, 0.809466, 0.683003,  0.722282, -1., 0.757329, 0.528404, 0.547243, 0.912726,  0.996799, -1., 0.676358, 0.789287, 0.304899, 0.859053,  0.733538, -1., 0.742325, 0.967502, 0.829836, 0.00964125,  0.789918, -1., -1., 0.179224, 0.836547, 0.334354, 0.316588,  0.951593, -1., 0.775051, 0.524888, 0.633585, 0.229311,  0.0177225, -1., 0.996484, 0.086342, 0.316585, 0.0209234,  0.320126, -1., 0.297055, 0.0116859, 0.16187, 0.586587,  0.554731, -1., -1., 0.0441843, 0.332034, 0.576946, 0.764813,  0.864961, -1., 0.495486, 0.242592, 0.448224, 0.913367,  0.720435, -1., 0.717703, 0.814639, 0.684056, 0.702713,  0.721219, -1., 0.728297, 0.367471, 0.681789, 0.401093,  0.431242, -1.}, 0.355786}, {0.00252397, 0.000177446, 0.103124,  0.00482737,  0.889347}}, {{{0.519919, 0.814506, 0.876511, 0.311601,  0.187886}, {{2, 3}, {0, 4}, {2, 4}, {3, 1}, {3, 4}, {0, 4}, {2,  3}, {1, 2}}, {-1., 0.23756, 0.111699, 0.446641, 0.692399,  0.994968, -1., 0.663474, 0.533273, 0.971964, 0.277265,  0.848835, -1., 0.849217, 0.269251, 0.556046, 0.120538,  0.481746, -1., 0.587462, 0.154952, 0.689297, 0.12596,  0.0675429, -1., -1., 0.340446, 0.812786, 0.814359, 0.879657,  0.102886, -1., 0.701087, 0.367718, 0.187258, 0.107918,  0.0376129, -1., 0.834445, 0.215294, 0.830653, 0.188778,  0.985228, -1., 0.946043, 0.274607, 0.0682393, 0.503482,  0.358581, -1., -1., 0.119655, 0.378943, 0.377522, 0.291038,  0.779209, -1., 0.566157, 0.563163, 0.41138, 0.676322,  0.86507, -1., 0.195445, 0.224122, 0.568404, 0.827457, 0.361, -1.,  0.0088278, 0.737751, 0.63868, 0.375772, 0.0627848, -1., -1.,  0.463144, 0.57044, 0.87229, 0.704204, 0.343488, -1., 0.191498,  0.494768, 0.413166, 0.56428, 0.625341, -1., 0.931606, 0.00178576,  0.887957, 0.760271, 0.736161, -1., 0.777664, 0.319553, 0.932813,  0.375161, 0.768836, -1.}, 0.581802}, {0.0000155295, 0.000142334,  0.0172814, 0.00152315,  0.981038}}, {{{0.294134, 0.999389, 0.706051, 0.118658,  0.723694}, {{3, 1}, {2, 4}, {3, 0}, {2, 1}, {0, 1}, {0, 1}, {3,  3}, {2, 4}}, {-1., 0.127099, 0.00184686, 0.77517, 0.532196,  0.632331, -1., 0.588681, 0.21089, 0.906856, 0.700725,  0.586895, -1., 0.322933, 0.146585, 0.964564, 0.809231,  0.00338039, -1., 0.213772, 0.589403, 0.0403955, 0.421579,  0.919638, -1., -1., 0.590014, 0.334345, 0.30292, 0.195944,  0.462915, -1., 0.332498, 0.527751, 0.663748, 0.830584,  0.743817, -1., 0.31686, 0.756893, 0.129859, 0.156922,  0.993927, -1., 0.610307, 0.165295, 0.347691, 0.990547,  0.396536, -1., -1., 0.575892, 0.307295, 0.568968, 0.476898,  0.985879, -1., 0.972951, 0.266048, 0.280954, 0.522964,  0.640453, -1., 0.738297, 0.617206, 0.69238, 0.896636,  0.421437, -1., 0.860313, 0.562521, 0.739714, 0.42751,  0.250006, -1., -1., 0.397226, 0.392023, 0.436963, 0.85347,  0.821333, -1., 0.0847277, 0.867995, 0.376572, 0.835455,  0.111777, -1., 0.601947, 0.0956181, 0.312491, 0.471323,  0.86365, -1., 0.478412, 0.620111, 0.574687, 0.442213,  0.618099, -1.}, 0.0575904}, {0.00563772, 0.944461, 0.000640023,  0.000272739,  0.0489887}}, {{{0.834973, 0.0147036, 0.368093, 0.660365,  0.44295}, {{3, 1}, {0, 3}, {2, 4}, {0, 3}, {0, 4}, {3, 0}, {3,  3}, {3, 4}}, {-1., 0.577741, 0.514624, 0.839031, 0.358222,  0.709746, -1., 0.138052, 0.00357645, 0.246445, 0.107799,  0.0424335, -1., 0.691085, 0.775122, 0.244149, 0.564021,  0.0709743, -1., 0.200435, 0.801936, 0.945922, 0.0133839,  0.365462, -1., -1., 0.787232, 0.577829, 0.353019, 0.922512,  0.209491, -1., 0.063205, 0.513988, 0.56429, 0.499745,  0.925153, -1., 0.510411, 0.317845, 0.391947, 0.88272,  0.819326, -1., 0.542722, 0.147798, 0.318699, 0.748352,  0.342288, -1., -1., 0.345862, 0.372777, 0.734968, 0.976826,  0.55863, -1., 0.794948, 0.381949, 0.0543134, 0.349138,  0.731743, -1., 0.867961, 0.490023, 0.849393, 0.80659,  0.357549, -1., 0.172179, 0.457446, 0.92387, 0.538223,  0.629456, -1., -1., 0.309649, 0.605171, 0.789872, 0.287169,  0.963787, -1., 0.232394, 0.054904, 0.310343, 0.405157,  0.437446, -1., 0.672956, 0.25603, 0.0560186, 0.705703,  0.804995, -1., 0.766007, 0.206626, 0.899113, 0.447446,  0.593828, -1.}, 0.749179}, {0.000103017, 0.000317183,  3.9011e-6, 0.78154,  0.218036}}, {{{0.975243, 0.909222, 0.964371, 0.439531,  0.370072}, {{3, 0}, {0, 0}, {3, 2}, {1, 2}, {0, 0}, {0, 1}, {0,  0}, {0, 3}}, {-1., 0.119351, 0.677202, 0.475744, 0.137678,  0.0644465, -1., 0.366859, 0.0705872, 0.700232, 0.391491,  0.110829, -1., 0.0145686, 0.994529, 0.586496, 0.344823,  0.807943, -1., 0.0954162, 0.139051, 0.750995, 0.0587633,  0.120173, -1., -1., 0.229828, 0.786624, 0.619232, 0.750101,  0.110478, -1., 0.109421, 0.143488, 0.612423, 0.0460312,  0.742562, -1., 0.0729011, 0.912191, 0.65454, 0.631733,  0.0583325, -1., 0.917662, 0.068044, 0.28691, 0.25039,  0.822246, -1., -1., 0.928993, 0.535915, 0.191626, 0.702073,  0.699165, -1., 0.749291, 0.572394, 0.951972, 0.588687,  0.63987, -1., 0.428906, 0.339549, 0.542656, 0.897308,  0.356005, -1., 0.427357, 0.888116, 0.265575, 0.297672,  0.509695, -1., -1., 0.820072, 0.978665, 0.0472824, 0.68745,  0.891079, -1., 0.44275, 0.855656, 0.985377, 0.191914,  0.693459, -1., 0.283262, 0.0334053, 0.603226, 0.0535889,  0.854356, -1., 0.693857, 0.0605702, 0.156281, 0.498352,  0.266499, -1.}, 0.172454}, {0.81948, 0.0402309, 0.135739,  0.00276423, 0.00178548}}}; 

  for (unsigned int i = 0; i < sizeof (data) / sizeof (data[0]); ++i)
    {
      Nominal nominal (data[i].inputs.alpha, 0, 0, 4, 5);
      std::vector<float> pz (5);
      std::vector<float> logpriorz (5);
      std::vector<Rating> ratings (8);

      for (unsigned int j = 0; j < 5; ++j) 
        { 
          logpriorz[j] = log (data[i].inputs.priorz[j]); 
          pz[j] = drand48 ();

          for (unsigned int k = 0; k < 4; ++k)
            {
              data[i].inputs.alpha[k * 5 * 5 + j * 5 + j] = drand48 ();
            }
        }

      for (unsigned int j = 0; j < 8; ++j) 
        {
          ratings[j] = data[i].inputs.ratings[j];
        }

      estep (nominal,
             data[i].inputs.logbeta,
             pz,
             logpriorz,
             ratings);

      for (unsigned int j = 0; j < 5; ++j)
        {
          assert (fabs (pz[j] - data[i].pz[j]) <=
                  1e-4 * (1 + fabs (pz[j]) + fabs (data[i].pz[j])) ||
                  (std::cerr << j << ": " << pz[j] << " ?= " << data[i].pz[j] << std::endl,
                   0));
        }
    }
}

}

int 
main (void) 
{
  srand48 (69);

  test_estep ();

  return 0;
}
